package com.jhm.ai.controller;

import com.alibaba.fastjson2.JSONObject;
import com.jhm.ai.config.AiConfig;
import com.jhm.ai.domain.AiChatSession;
import com.jhm.ai.domain.ChatRequest;
import com.jhm.ai.service.IAiChatMessageService;
import com.jhm.ai.service.IAiChatSessionService;
import com.jhm.ai.util.HttpConnectionManager;
import com.jhm.common.core.controller.BaseController;
import com.jhm.common.core.domain.AjaxResult;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.net.HttpURLConnection;
import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping("/ai/chat")
public class AiChatController extends BaseController {
    @Autowired
    private IAiChatSessionService sessionService;

    @Autowired
    private IAiChatMessageService messageService;

    @Autowired
    private AiConfig aiConfig;

    /**
     * 获取请求ID
     */
    @GetMapping("/requestId")
    public AjaxResult getRequestId() {
        String requestId = HttpConnectionManager.generateRequestId();
        return AjaxResult.success("获取成功", requestId);
    }

    /**
     * 创建新会话
     */
    @PostMapping("/session")
    public AjaxResult createSession(@RequestBody AiChatSession session) {
        session.setCreateBy(getUsername());
        return toAjax(sessionService.insertAiChatSession(session));
    }

    /**
     * 获取当前用户的会话列表
     */
    @GetMapping("/session/list")
    public AjaxResult listSessions() {
        AiChatSession session = new AiChatSession();
        session.setCreateBy(getUsername());
        return AjaxResult.success(sessionService.selectAiChatSessionList(session));
    }

    /**
     * 修改会话标题
     */
    @PutMapping("/session")
    public AjaxResult updateSession(@RequestBody AiChatSession session) {
        return toAjax(sessionService.updateAiChatSession(session));
    }

    /**
     * 获取会话的消息列表
     */
    @GetMapping("/message/{sessionId}")
    public AjaxResult getSessionMessages(@PathVariable("sessionId") Long sessionId) {
        return AjaxResult.success(messageService.selectMessagesBySessionId(sessionId));
    }

    /**
     * 删除会话
     */
    @DeleteMapping("/session/{sessionId}")
    public AjaxResult removeSession(@PathVariable Long sessionId) {
        return toAjax(sessionService.deleteAiChatSessionById(sessionId));
    }

    /**
     * SSE聊天接口
     */
    @PostMapping("/sse")
    public void sseChat(@RequestBody ChatRequest chatRequest, HttpServletResponse response) throws IOException {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");
        
        PrintWriter writer = response.getWriter();
        
        try {
            messageService.handleSseChat(chatRequest, getUsername(), chatRequest.getRequestId(), writer);
        } catch (Exception e) {
            Map<String, String> errorMap = new HashMap<>();
            errorMap.put("error", "服务器处理异常，请稍后重试");
            errorMap.put("type", "error");
            writer.write("data: " + JSONObject.toJSONString(errorMap) + "\n\n");
            writer.flush();
        } finally {
            if (chatRequest.getRequestId() != null) {
                HttpConnectionManager.removeConnection(chatRequest.getRequestId());
            }
            writer.close();
        }
    }

    /**
     * 停止SSE生成
     */
    @PostMapping("/stop/{requestId}")
    public AjaxResult stopGeneration(@PathVariable String requestId) {
        HttpURLConnection conn = HttpConnectionManager.getConnection(requestId);
        if (conn != null) {
            try {
                conn.disconnect();
                HttpConnectionManager.removeConnection(requestId);
                return AjaxResult.success("停止成功");
            } catch (Exception e) {
                return AjaxResult.error("停止失败");
            }
        }
        return AjaxResult.error("未找到对应的请求");
    }
}
